from flask import Flask, render_template, request, jsonify
from sparkai.llm.llm import ChatSparkLLM, ChunkPrintHandler
from sparkai.core.messages import ChatMessage
import os
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# 从环境变量中读取星火认知大模型相关信息
SPARKAI_URL = 'wss://spark-api.xf-yun.com/v3.5/chat'
SPARKAI_APP_ID = os.getenv('SPARKAI_APP_ID')
SPARKAI_API_SECRET = os.getenv('SPARKAI_API_SECRET')
SPARKAI_API_KEY = os.getenv('SPARKAI_API_KEY')
SPARKAI_DOMAIN = 'generalv3.5'

app = Flask(__name__)

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/chat', methods=['POST'])
def chat():
    try:
        # 从请求中获取用户输入
        data = request.get_json()
        user_input = data.get('user_input', '')

        if not user_input:
            return jsonify({'error': '用户输入不能为空'}), 400

        # 创建消息对象
        messages = [ChatMessage(role="user", content=user_input)]

        # 初始化星火认知大模型
        spark = ChatSparkLLM(
            SPARKAI_URL = 'wss://spark-api.xf-yun.com/v3.5/chat'
            SPARKAI_APP_ID = '982c0054'
            SPARKAI_API_SECRET = 'NzRjMTdhNWUyYmU4MTNiNGQxZDgzMWQ0'
            SPARKAI_API_KEY = 'f6a31ec70344fd3350694509a8fec7e3'
            SPARKAI_DOMAIN = 'generalv3.5'
            streaming=False,
        )

        # 使用星火认知大模型生成响应
        handler = ChunkPrintHandler()
        response = spark.generate([messages], callbacks=[handler])

        # 处理响应
        if response and response[0].text:
            text_response = response[0].text
        else:
            text_response = "对不起，我无法理解您的问题。"

        # 返回响应给前端
        return jsonify({'response': text_response})

    except Exception as e:
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    # 初始化星火认知大模型以确保配置正确
    try:
        spark = ChatSparkLLM(
            spark_api_url=SPARKAI_URL,
            spark_app_id=SPARKAI_APP_ID,
            spark_api_key=SPARKAI_API_KEY,
            spark_api_secret=SPARKAI_API_SECRET,
            spark_llm_domain=SPARKAI_DOMAIN,
            streaming=False,
        )
        messages = [ChatMessage(role="user", content='你好呀')]
        handler = ChunkPrintHandler()
        a = spark.generate([messages], callbacks=[handler])
        print(a)
    except Exception as e:
        print(f"初始化模型时出错: {e}")

    # 运行 Flask 应用
    app.run(debug=True)